#include "lerp_tiling.h"
#include "register/op_def_registry.h"
#include "tiling/platform/platform_ascendc.h"

namespace optiling {
const uint32_t BLOCK_SIZE = 32;
static ge::graphStatus TilingFunc(gert::TilingContext* context)
{

  LerpTilingData tiling;
    uint32_t sizeofdatatype;
    uint32_t datatypeclass;
    uint32_t totalLengthAligned;
    uint32_t ALIGN_NUM;
    uint32_t blockLength;
    uint32_t tileLength;
    uint32_t lasttileLength;
    uint64_t ub_size;
    uint32_t tile_num;
    uint32_t startLength;
    uint32_t endLength;
    uint32_t weightLength;
    uint32_t formerLength = 0;
    uint32_t formertileLength = 0;
    uint32_t formerlasttileLength = 0;
    uint32_t broadcast = 0;
    uint32_t totalLength = context->GetInputTensor(0)->GetShapeSize();
    uint32_t endlastDim;
    uint32_t startlastDim;
    auto dt = context->GetInputTensor(0)->GetDataType();
    if (dt == 1)
    { // float16
      sizeofdatatype = 2;
      datatypeclass = 0;

    }
    else if (dt == 0)
    { // float
      sizeofdatatype = 4;
      datatypeclass = 1;
    }
    ALIGN_NUM = BLOCK_SIZE / sizeofdatatype;

    // 平台信息获取
    auto ascendcPlatform = platform_ascendc::PlatformAscendC(context->GetPlatformInfo());
    auto socVersion = ascendcPlatform.GetSocVersion();
    ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ub_size);
    auto aivNum = ascendcPlatform.GetCoreNumAiv();
    uint32_t ub_block_num = 4;
    if (totalLength >= 1024*1024){
	    ub_block_num = 256;
    }
    // input_size

    auto endShape= context->GetInputTensor(1)->GetOriginShape();
    auto startShape = context->GetInputTensor(0)->GetOriginShape();
    startLength = context->GetInputTensor(0)->GetShapeSize();
    endLength = context->GetInputTensor(1)->GetShapeSize();
    endlastDim = endShape.GetDim(endShape.GetDimNum() - 1);
    startlastDim = startShape.GetDim(startShape.GetDimNum() - 1);
    weightLength = context->GetInputTensor(2)->GetShapeSize();
    uint32_t broadcastend = startLength / endLength;
    if (broadcastend > 1){
	    broadcast = 1;
	    if (endlastDim == 1 && endLength > 1){ broadcast = 2;}
    }

    printf("endlastDim %d \n",endlastDim);
    // 32Aligned
    if (totalLength % ALIGN_NUM != 0)
    {
      totalLengthAligned =
          ((totalLength + ALIGN_NUM - 1) / ALIGN_NUM) * ALIGN_NUM;
    }
    else
    {
      totalLengthAligned = totalLength;
    }

    // SetBlockDim()
    if (totalLengthAligned <= ub_block_num * ALIGN_NUM)
    { // shape较小，用单核
      context->SetBlockDim(1);
    }
    else
    {
      if (((totalLengthAligned / ALIGN_NUM) % ub_block_num) ==
          0)
      { // 可以核间均分
        if ((totalLengthAligned / ALIGN_NUM / ub_block_num) <=
            aivNum)
        { // 且计算出均分后的核数小于当前aicore数量，按计算值
          context->SetBlockDim(totalLengthAligned / ALIGN_NUM / ub_block_num);
        }
        else
        {
          // ... 按照aivNum切分
          context->SetBlockDim(aivNum);
        }
      }
      else
      { // 核间不能均分
        if (((totalLengthAligned / ALIGN_NUM / ub_block_num) + 1) <=
            aivNum)
        { // 且计算出均分后的核数小于当前aicore数量，按计算值
          context->SetBlockDim((totalLengthAligned / ALIGN_NUM / ub_block_num) + 1);
        }
        else
        {
          // ... 按照aivNum切分
          context->SetBlockDim(aivNum);
        }
      }
    }

    // infinite value
    auto block_dim = context->GetBlockDim();
    
    if ((totalLengthAligned / ALIGN_NUM) % block_dim == 0)
    { // 核间可均分
      blockLength = totalLengthAligned / block_dim;
      tile_num = blockLength / ALIGN_NUM / ub_block_num; // 分块数量
      if ((blockLength / ALIGN_NUM) % ub_block_num == 0 ||
          tile_num == 0)
      { 
        if (tile_num == 0)
        {
          tile_num = 1;
        }
        if (blockLength < ub_block_num) {
        tileLength = ((blockLength / ALIGN_NUM) + 1) / 2 * 2 * ALIGN_NUM;
        lasttileLength = tileLength;
      } else {
        tileLength = ub_block_num * ALIGN_NUM;
        lasttileLength = tileLength;
      }
      }
      else
      { // 满足32字节对齐，核内不能均分， 存在尾块
        tile_num = tile_num + 1;
        tileLength = ub_block_num * ALIGN_NUM;
        lasttileLength = blockLength - (tile_num - 1) * tileLength;
      }
      // 调试数据打印
      printf("----------------核间均分--------------\n");
      printf("blockDim: %d \n", block_dim);
      printf("blockLength: %d \n", blockLength);
      printf("tileNum: %d \n", tile_num);
      printf("ub_block_num %d \n", ub_block_num);
      printf("ALIGN_NUM: %d \n", ALIGN_NUM);
      printf("tileLength: %d \n", tileLength);
      printf("lasttileLength: %d \n", lasttileLength);
      printf("startLength: %d \n", startLength);
      printf("endLength: %d \n", endLength);
      printf("weightLength: %d \n", weightLength);
      printf("datatypeclass: %d \n",datatypeclass);
      // tiling信息填充
      context->SetTilingKey(1);
      tiling.set_blockLength(blockLength);
      tiling.set_tileNum(tile_num);
      tiling.set_tileLength(tileLength);
      tiling.set_lasttileLength(lasttileLength);
      tiling.set_startLength(startLength);
      tiling.set_endLength(endLength);
      tiling.set_datatypeclass(datatypeclass);
      tiling.set_weightLength(weightLength);
      tiling.set_broadcast(broadcast);
      tiling.set_startlastDim(startlastDim);
      tiling.SaveToBuffer(context->GetRawTilingData()->GetData(),
                          context->GetRawTilingData()->GetCapacity());
      context->GetRawTilingData()->SetDataSize(tiling.GetDataSize());
      size_t *currentWorkspace = context->GetWorkspaceSizes(1);
      currentWorkspace[0] = 0;
      return ge::GRAPH_SUCCESS;
    }
    else
    {
      // 计算大小快的数量及数据量
      uint32_t formerNum = (totalLengthAligned / ALIGN_NUM) % block_dim;
      uint32_t tailNum = block_dim - formerNum;
      uint32_t formerLength =
          (((totalLengthAligned + block_dim - 1) / block_dim + ALIGN_NUM - 1) /
           ALIGN_NUM) *
          ALIGN_NUM;
      uint32_t tailLength = (totalLengthAligned / block_dim / ALIGN_NUM) * ALIGN_NUM;

      // former 情况
      uint32_t former_tile_num = formerLength / ALIGN_NUM / ub_block_num;
      if ((formerLength / ALIGN_NUM) % ub_block_num == 0 ||
          former_tile_num == 0)
      { // 核内均分
        if (former_tile_num == 0)
        {
          former_tile_num = 1;
        }
      if (formerLength < ub_block_num * ALIGN_NUM) {
        formertileLength = ((formerLength / ALIGN_NUM) + 1) / 2 * 2 * ALIGN_NUM;
        formerlasttileLength = formertileLength;
      } else {
        formertileLength = ub_block_num * ALIGN_NUM;
        formerlasttileLength = formertileLength;
      }	
      }
      else
      {
        former_tile_num = former_tile_num + 1;
        formertileLength = ub_block_num * ALIGN_NUM;
        formerlasttileLength =
            (formerLength - (former_tile_num - 1) * formertileLength);
      }

      // tail情况
      uint32_t tail_tile_num = tailLength / ALIGN_NUM / ub_block_num;
      uint32_t tailtileLength;
      uint32_t taillasttileLength;
      if ((tailLength / ALIGN_NUM) % ub_block_num == 0 ||
          tail_tile_num == 0)
      { // 核内可以均分
        if (tail_tile_num == 0)
        {
          tail_tile_num = 1;
        }
      if (tailLength < (ub_block_num * ALIGN_NUM)) {
        tailtileLength = ((tailLength / ALIGN_NUM) + 1) / 2 * 2 * ALIGN_NUM;
        taillasttileLength = tailtileLength;
      } else {
        tailtileLength = ub_block_num * ALIGN_NUM;
        taillasttileLength = tailtileLength;
      }
      }
      else
      { // 核内不均分
        tail_tile_num = tail_tile_num + 1;
        tailtileLength = ub_block_num * ALIGN_NUM;
        taillasttileLength = (tailLength - (tail_tile_num - 1) * tailtileLength);
      }
      // 打印调试信息
      printf("---------------核间非均分------------- \n");
      printf("formerNum: %d \n", formerNum);
      printf("tailNum: %d \n", tailNum);
      printf("formerLength: %d \n", formerLength);
      printf("tailLength: %d \n", tailLength);
      printf("formertileNum: %d \n", former_tile_num);
      printf("tailtileNum: %d \n", tail_tile_num);
      printf("blockDim: %d \n", block_dim);
      printf("totalLengthAligned: %d \n", totalLengthAligned);
      printf("formerlasttileLength: %d \n", formerlasttileLength);
      printf("taillasttileLength: %d \n", taillasttileLength);
      // tiling填充
      tiling.set_startLength(startLength);
      tiling.set_endLength(endLength);
      tiling.set_weightLength(weightLength);
      tiling.set_formerNum(formerNum);
      tiling.set_formerLength(formerLength);
      tiling.set_formertileNum(former_tile_num);
      tiling.set_formertileLength(formertileLength);
      tiling.set_formerlasttileLength(formerlasttileLength);
      tiling.set_tailNum(tailNum);
      tiling.set_datatypeclass(datatypeclass);
      tiling.set_tailLength(tailLength);
      tiling.set_tailtileNum(tail_tile_num);
      tiling.set_tailtileLength(tailtileLength);
      tiling.set_taillasttileLength(taillasttileLength);
      tiling.set_broadcast(broadcast);
      tiling.set_startlastDim(startlastDim);
      context->SetTilingKey(2);
      tiling.SaveToBuffer(context->GetRawTilingData()->GetData(),
                          context->GetRawTilingData()->GetCapacity());
      context->GetRawTilingData()->SetDataSize(tiling.GetDataSize());
      size_t *currentWorkspace = context->GetWorkspaceSizes(1);
      currentWorkspace[0] = 0;
      return ge::GRAPH_SUCCESS;
    }
  }
}


namespace ge {
static ge::graphStatus InferShape(gert::InferShapeContext* context)
{
    const gert::Shape* x1_shape = context->GetInputShape(0);
    gert::Shape* y_shape = context->GetOutputShape(0);
    *y_shape = *x1_shape;
    return GRAPH_SUCCESS;
}
}


namespace ops {
class Lerp : public OpDef {
public:
    explicit Lerp(const char* name) : OpDef(name)
    {
        this->Input("start")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT16, ge::DT_FLOAT})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});
        this->Input("end")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT16, ge::DT_FLOAT})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});
        this->Input("weight")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT16, ge::DT_FLOAT})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});
        this->Output("y")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT16, ge::DT_FLOAT})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});

        this->SetInferShape(ge::InferShape);

        this->AICore()
            .SetTiling(optiling::TilingFunc);
        this->AICore().AddConfig("ascend310b");

    }
};

OP_ADD(Lerp);
}

